"""模型注册表模块"""

from __future__ import annotations

from pathlib import Path
from typing import Dict, Any

# 模型注册表：用于存放已知的标准模型信息
# 每个模型包含模型名、ONNX 路径、输入/输出描述、任务类型、opset 版本、SHA256 哈希
import json

def _load_registry_from_json() -> Dict[str, Dict[str, Any]]:
    REGISTRY_JSON_PATH = Path(__file__).parent / "model_registry.json"
    if not REGISTRY_JSON_PATH.exists():
        raise FileNotFoundError(f"模型注册表文件不存在: {REGISTRY_JSON_PATH}")
    with open(REGISTRY_JSON_PATH, "r", encoding="utf-8") as f:
        return json.load(f)

_REGISTRY: Dict[str, Dict[str, Any]] = _load_registry_from_json()

def get_model_spec(model_or_path: str) -> Dict[str, Any]:
    """
    获取模型规范（字典结构）。
    支持两种输入：
      1. 已注册的模型名（如 "resnet50"）
      2. 指定的 ONNX 路径（如 "path/to/my_model.onnx"），要求文件存在
    返回：
      - 包含模型名、ONNX 文件路径、输入信息、任务类型等
    若找不到指定模型或路径不存在会抛出 FileNotFoundError
    """
    # 如果是注册表中已知的模型名，则返回注册信息
    if model_or_path in _REGISTRY:
        spec = _REGISTRY[model_or_path].copy()
        # 确保路径存在
        model_path = Path(spec["source_onnx"])
        if not model_path.exists():
            raise FileNotFoundError(
                f"Model file not found: {model_path}\n"
                f"Please download it first using: python scripts/download_models.py --name {model_or_path}"
            )
        return spec
        
    # 如果参数为本地 ONNX 路径，且文件存在，则动态构建规范
    p = Path(model_or_path)
    if p.suffix.lower() == ".onnx" and p.exists():
        return {
            "name": p.stem,
            "source_onnx": str(p.as_posix()),
            "task": "unknown",
            "opset": None,
            "sha256": None,
            "inputs": {
                "input": {
                    "shape": [1, 3, 224, 224],
                    "dtype": "float32",
                    "layout": "NCHW",
                    "batched": True
                }
            },
            "outputs": ["output"],
        }
    
    # 未知模型或路径
    raise FileNotFoundError(
        f"Unknown model or path: {model_or_path}\n"
        f"Available models: {list(_REGISTRY.keys())}\n"
        f"Or provide a valid ONNX file path."
    )

def list_models():
    """
    返回所有已注册模型的名称列表
    """
    return list(_REGISTRY.keys())

